import json
from fastapi import FastAPI, Form, File, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from io import BytesIO
from .templates import form_page
from .email_sender import send_email
from .utils import load_email_data, personalize_body
import asyncio
import os

app = FastAPI()

# WebSocket 连接的列表
connected_clients = []
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许来自这个地址的请求
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.websocket("/ws/send-email")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    connected_clients.append(websocket)
    try:
        while True:
            await asyncio.sleep(1)  # 保持连接，避免关闭
    except WebSocketDisconnect:
        connected_clients.remove(websocket)

# 修改 `process_email` 函数
@app.post("/send-email")
async def process_email(
        smtp_user: str = Form(...),
        smtp_password: str = Form(...),
        subject: str = Form(...),
        body: str = Form(...),
        file: UploadFile = File(...),  # 用于上传包含电子邮件和文件名的 Excel 文件
        attachments: list[UploadFile] = File(...)  # 接受多个附件
):
    try:
        # 将上传的文件内容读取为字节流并传递给 load_email_data
        file_bytes = await file.read()  # 读取文件内容
        file_stream = BytesIO(file_bytes)  # 转换为字节流

        # 调用 load_email_data 并传入字节流
        email_list, file_names = load_email_data(file_stream)
    except ValueError as e:
        return {"message": str(e)}

    success_list = []
    failed_list = []
    total_emails = len(email_list)

    # 处理上传的附件
    for attachment in attachments:
        attachment_bytes = await attachment.read()
        file_path = os.path.join("attachments", attachment.filename)  # 设置附件保存路径
        with open(file_path, "wb") as f:
            f.write(attachment_bytes)

    for index, (email, file_name) in enumerate(zip(email_list, file_names)):
        file_path = f"attachments/{file_name}"
        if os.path.exists(file_path):  # 检查文件是否存在
            personalized_body = personalize_body(body, email)
            try:
                success, error = send_email(smtp_user, smtp_password, email, subject, personalized_body, file_path)
                if success:
                    success_list.append(email)
                else:
                    failed_list.append(f"{email} (发送失败: {error})")
            except Exception as e:
                failed_list.append(f"{email} (发送时发生错误: {str(e)})")
        else:
            failed_list.append(f"{email} (未找到附件: {file_name})")

        # 发送进度更新
        progress_percentage = int(((index + 1) / total_emails) * 100)
        await send_progress_update(progress_percentage)

    if failed_list:
        return {
            "message": "Some emails failed to send.",
            "successfully_sent": success_list,
            "failed_to_send": failed_list
        }

    return {
        "message": "Emails sent successfully!",
        "successfully_sent": success_list
    }

async def send_progress_update(progress: int):
    # 向所有连接的 WebSocket 客户端发送进度更新
    for client in connected_clients:
        try:
            await client.send_text(json.dumps({"progress": progress}))  # 使用 JSON 格式发送进度
        except Exception as e:
            # 处理连接关闭的情况
            print(f"Connection closed: {e}")
            connected_clients.remove(client)
